Docs: Improve SFT/RL user experience #2794
Open
+25
−23
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
hengtaoguo
force-pushed
the
hengtaoguo-grpo
branch
2 times, most recently
from
8629e8b to
5ae647f
Compare
hengtaoguo
requested review from
A9isha,
RissyRan,
bvandermoon,
gagika,
gobbleturk,
jacoguzo,
jiangjy1982,
richjames0,
shralex and
vipannalla
as code owners
hengtaoguo
force-pushed
the
hengtaoguo-grpo
branch
from
c836743 to
66c8566
Compare
NuojCheng
approved these changes
Dec 10, 2025
hengtaoguo
force-pushed
the
hengtaoguo-grpo
branch
from
d971a0b to
9b41868
Compare
hengtaoguo
force-pushed
the
hengtaoguo-grpo
branch
from
b87dc05 to
4f3a8cd
Compare
|
|
||
| ## Create virtual environment and Install MaxText dependencies | ||
| If you have already completed the [MaxText installation](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/install_maxtext.md), you can skip to the next section for post-training dependencies installations. Otherwise, please install `MaxText` using the following commands before proceeding. | ||
| If you have already completed the [MaxText installation](../../install_maxtext.md), you can skip to the next section for post-training dependencies installations. Otherwise, please install `MaxText` using the following commands before proceeding. |
Collaborator
There was a problem hiding this comment.
Why do we need to change the link here?
|
|
||
| export RUN_NAME=<name for this run> # e.g., $(date +%Y-%m-%d-%H-%M-%S) | ||
| export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/0/items | ||
| export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/0/items # Actual checkpoint saved with an extra /0/items path suffix |
Collaborator
There was a problem hiding this comment.
This doesn't look right if user has the checkpoint in a GCS. We can remove this env variable from here and move this to next section, similar to https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/sft.html#get-your-model-checkpoint.
| The overview of what this run will do is as follows: | ||
|
|
||
| 1. We load a policy model and a reference model. Both are copies of `Llama3.1-8b-Instruct`. | ||
| 1. We load a policy model and a reference model. Both are copies of the model checkpoint you specified (e.g., `Llama3.1-8b-Instruct`). |
Collaborator
There was a problem hiding this comment.
Can you do the same at line 128?
|
|
||
| ## 2. Install XPK | ||
| Install XPK by following the instructions in the [official documentation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation-via-pip). | ||
| Install XPK by following the instructions in the [official documentation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation-via-pip). We also provide a quick guide for XPK installation and usage [here](https://maxtext.readthedocs.io/en/latest/run_maxtext/run_maxtext_via_xpk.html). |
Collaborator
There was a problem hiding this comment.
nit: We also provide a quick guide for XPK installation here
That XPK documentation mainly talks about pre-training. Pointing users to XPK documentation at this point might create some confusion. Can we explicitly say just follow the instruction in that guide for XPK installation & Prerequisite and continue on the current doc for post-training?
| ## Submit your RL workload via Pathways | ||
|
|
||
| Please create a pathways ready GKE cluster as described [here](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster), and you can submit the `train_rl.py` script via [XPK](https://github.com/AI-Hypercomputer/xpk). | ||
| Please create a pathways ready GKE cluster as described [here](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster), and you can submit the `train_rl.py` script via [XPK](https://github.com/AI-Hypercomputer/xpk). We also provide a quick guide for XPK installation and usage [here](../../run_maxtext/run_maxtext_via_xpk.md). |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Reduce user friction in SFT/RL and fix broken links.
b/463394566
b/463409639
b/463409807
b/463396352
b/463393644
Tests
N/A
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.